package sunyu.tools.kafka;

import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import javax.annotation.PreDestroy;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Future;

/**
 * @author 孙宇
 */
@Component
public class KafkaTools {

    public static final String defaultName = "main";
    private final Logger logger = LoggerFactory.getLogger(this.getClass());
    private ConcurrentHashMap<String, Producer> producerMap = new ConcurrentHashMap<>();
    private ConcurrentHashMap<String, KafkaConsumerTemplate> consumerMap = new ConcurrentHashMap<>();

    public void addProducer(String name,
                            String servers,
                            String keySerializer,
                            String valueSerializer) {
        logger.info("添加生产者：{}", name);
        if (producerMap.containsKey(name)) {
            logger.error("{}已存在!", name);
        } else {
            Properties props = new Properties();
            props.put("bootstrap.servers", servers);
            props.put("key.serializer", keySerializer);
            props.put("value.serializer", valueSerializer);
            /*props.put("acks", "all");
            props.put("retries ", 1);
            props.put("buffer.memory", 33554432);
            */
            Producer<String, String> producer = new KafkaProducer<>(props);
            producerMap.put(name, producer);
        }
    }

    public Producer getProducer(String name) {
        return producerMap.get(name);
    }

    public Future send(String producerName,
                       ProducerRecord record) {
        return producerMap.get(producerName).send(record);
    }

    public void addConsumer(String name,
                            KafkaConsumerTemplate consumer) {
        logger.info("添加消费者：{}", name);
        if (consumerMap.containsKey(name)) {
            logger.error("{}已存在!", name);
        } else {
            consumerMap.put(name, consumer);
        }
    }

    public KafkaConsumerTemplate getConsumer(String name) {
        return consumerMap.get(name);
    }

    @PreDestroy
    public void destroy() {
        logger.info("销毁所有生产者和消费者开始");
        for (Map.Entry<String, Producer> entry : producerMap.entrySet()) {
            entry.getValue().close();
        }

        for (Map.Entry<String, KafkaConsumerTemplate> entry : consumerMap.entrySet()) {
            entry.getValue().setRunner(false);
        }
        logger.info("销毁所有生产者和消费者结束");
    }

}
